import numpy as np
import torch
from transformers import BertTokenizer, BertModel
from Classifier import Classifier
import warnings
warnings.filterwarnings("ignore")

model_name = '../bert-base-chinese'
model_path = '../bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
bert_model = BertModel.from_pretrained(model_path)
model = Classifier(bert_model)
# 加载最佳模型的权重
路径='saved_weights_外卖.pt'
model.load_state_dict(torch.load(路径))

while True:
    text=input("请输入对外卖的评论：")
    sent_id = tokenizer.encode(text,
                            add_special_tokens = True,
                            max_length = 100,
                            truncation = True,
                            pad_to_max_length = 'right')

    att_mask = [int(tok > 0) for tok in sent_id]
    sent_id = torch.tensor(sent_id)
    att_mask = torch.tensor(att_mask)
    sent_id = sent_id.unsqueeze(0)
    att_mask = att_mask.unsqueeze(0)
    preds = model(sent_id, att_mask)


    total_preds=[]
    total_preds.append(np.argmax(preds.detach().cpu().numpy(), axis=1))
    # print(np.concatenate(total_preds, axis=0))
    if 1 in np.concatenate(total_preds, axis=0):
        print("好评")
    else:
        print("差评")

